import numpy as np
import matplotlib.pyplot as plt


# 线性回归模型类
class LinearRegression:
    def __init__(self, learning_rate=0.01, n_iterations=1000):
        self.learning_rate = learning_rate
        self.n_iterations = n_iterations
        self.theta = None  # 模型参数
        self.cost_history = []  # 损失值历史记录

    def fit(self, X, y):
        m, n = X.shape  # 样本数量，特征数量
        self.theta = np.zeros(n + 1)  # 初始化参数，增加截距项
        X = np.c_[np.ones(m), X]  # 在特征矩阵中添加一列全为1的列

        # 梯度下降
        for _ in range(self.n_iterations):
            predictions = X.dot(self.theta)  # 计算预测值
            errors = predictions - y  # 计算误差
            gradient = (1 / m) * X.T.dot(errors)  # 计算梯度
            self.theta -= self.learning_rate * gradient  # 更新参数
            cost = (1 / (2 * m)) * np.sum(errors ** 2)  # 计算损失值
            self.cost_history.append(cost)  # 记录损失值

    def predict(self, X):
        X = np.c_[np.ones(X.shape[0]), X]  # 在特征矩阵中添加一列全为1的列
        return X.dot(self.theta)  # 返回预测值

    def get_cost_history(self):
        return self.cost_history  # 返回损失值历史记录


# 示例数据
np.random.seed(42)  # 设置随机种子以确保结果可重复
X = 2 * np.random.rand(100, 1)  # 生成100个样本，1个特征
y = 4 + 3 * X + np.random.randn(100, 1)  # 生成目标值，带有噪声
y = y.ravel()  # 将 y 转换为一维数组，形状为 (100,)

# 创建线性回归模型
model = LinearRegression(learning_rate=0.1, n_iterations=100)

# 训练模型
model.fit(X, y)

# 预测
X_new = np.array([[0], [2]])  # 新数据
predictions = model.predict(X_new)

# 输出预测结果
print("模型参数 theta:", model.theta)
print("预测值 (X_new = [[0], [2]]):", predictions)

# 绘制损失值变化曲线
plt.plot(model.get_cost_history())
plt.xlabel("迭代次数")
plt.ylabel("损失值")
plt.title("损失值变化曲线")
plt.show()

# 绘制数据点和回归线
plt.scatter(X, y, color="blue", label="数据点")
plt.plot(X_new, predictions, color="red", label="回归线")
plt.xlabel("特征 X")
plt.ylabel("目标值 y")
plt.title("线性回归拟合结果")
plt.legend()
plt.show()
